Add AMDGPU execution provider - #2165
Conversation
|
kunal-vaishnavi Baiju Meswani (@baijumeswani) could you help review this PR? |
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds MIGraphX (AMD GPU) execution provider integration to ONNX Runtime GenAI, including provider name normalization and runtime behaviors needed to avoid recompilation during prompt processing.
Changes:
- Register MIGraphX in the session-options dispatch table and add a MIGraphX EP implementation (V2 plugin path with V1 fallback).
- Normalize provider names so
"migraphx"and"MIGraphXExecutionProvider"map to"MIGraphX", and enable graph capture for MIGraphX. - Add “static input shape” prompt-time padding and update position/logits shapes to support padded prompt lengths.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
| src/models/session_options.cpp | Adds MIGraphX EP registration in the provider dispatch map. |
| src/models/position_inputs.h | Updates CreateAndInitializePositionIDs signature to accept actual seq_length. |
| src/models/position_inputs.cpp | Implements position id initialization w/ padded shapes and safe indexing using seq_length. |
| src/models/model.h | Adds State::prompt_gen_ flag to control prompt-time padding behavior. |
| src/models/model.cpp | Pads input_ids to max_length during prompt generation for static-shape EPs. |
| src/models/logits.cpp | Forces logits output shape to max_length during prompt generation. |
| src/migraphx/session_options.h | Declares MIGraphX EP append entrypoint. |
| src/migraphx/session_options.cpp | Implements MIGraphX EP append with V2 plugin-first and V1 fallback. |
| src/generators.h | Adds use_static_input_shapes to generator params. |
| src/generators.cpp | Toggles prompt_gen_ around prompt vs token-generation runs. |
| src/config.h | Declares NeedsStaticInputShapes. |
| src/config.cpp | Adds MIGraphX name normalization, enables graph capture for MIGraphX, and implements NeedsStaticInputShapes. |
| cmake/global_variables.cmake | Adds MIGraphX sources to the CMake glob. |
c2f4145 to
d563801
Compare
|
Hi kunal-vaishnavi. Apologies for the long delay getting back to you. In the meantime, we've been reworking the architecture of the AMD GPU integration, and your comments on this PR shaped that direction. This resolves the concerns you raised here:
We'll push the new-architecture updates to this same PR rather than opening a new one retitling/rescoping it to the umbrella EP as we go. Thanks for the patience and for the feedback that pushed us here. |
d563801 to
66e496a
Compare
|
Aditya Lohia (@aditya-dl) please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
|
Hi kunal-vaishnavi. The new-architecture updates are now pushed to this PR (three commits on top of latest main), and I've updated the description to match. The EP-specific if/else logic and the logits.cpp hackiness from the previous version are both gone. Whenever you have a chance, could you take another look? Happy to walk through any of it. Thanks again for the earlier feedback. |
…ice (review microsoft#2165) Resolve the AMDGPU device id from the filtered EP device instead of hardcoding 0. The EP keys its allocator on this id, so a hardcoded 0 pinned compute to device 0 regardless of the selected device. Correlate the host-accessible pool to the same id so pinned decode inputs live on the device that runs the model. Single-GPU resolves to id 0, unchanged.
…ice (review microsoft#2165) Resolve the AMDGPU device id from the filtered EP device instead of hardcoding 0. The EP keys its allocator on this id, so a hardcoded 0 pinned compute to device 0 regardless of the selected device. Correlate the host-accessible pool to the same id so pinned decode inputs live on the device that runs the model. Single-GPU resolves to id 0, unchanged.
Add AMDGPU execution provider (port of microsoft#2165 to rel-0.15.1)
Add the AMDGPU execution provider to ONNX Runtime GenAI. The AMDGPU EP resolves a profile to a concrete backend (MIGraphX or DML) at runtime; OGA drives it as a single DeviceType::AMDGPU. Provider naming: exposed as "amdgpu"; OGA also accepts "AMDGPUExecutionProvider" (the catalog form used by the AMD-shipped Windows ML EP MSIX) so test harnesses that match config strings against WinML-discovered names work without bypass hacks. Both normalize to "AMDGPU". GPU-resident KV cache: a GPU-resident DeviceInterface keeps the KV cache on the device (no per-token CPU-to-GPU roundtrip), with backend-agnostic opaque DeviceBuffer copies that dispatch to the active backend. Static-shape prefill: emit ep.migraphx.static_pad_* and hip_graph_enable session-config entries so the EP pads the prefill token axis and reuses a captured graph. DML ignores the migraphx-namespaced keys. Known limitations: - Beam search not supported (needs past_present_share_buffer=true, which requires num_beams=1)
Route the small decode inputs (input_ids/position_ids/attention_mask) through a host-accessible (CPU-writable, GPU-readable) allocator so the CPU updates them in place with no per-step copy. Resolved via GetSharedAllocator; KV cache and scoring stay on the default device interface, and the path falls back to default inputs if no host-accessible allocator is available. Single-GPU only for now (device_id 0).
Logits is GPU-written and CPU-read (the sampler), the opposite of the pinned decode inputs. On AMDGPU the inputs use a host-accessible allocator whose heap is not CPU-read-coherent, so reading logits from it returns stale data. Route logits to the CPU interface instead via a new p_logits_ member; only the decode inputs stay on the host-accessible allocator.
…ice (review microsoft#2165) Resolve the AMDGPU device id from the filtered EP device instead of hardcoding 0. The EP keys its allocator on this id, so a hardcoded 0 pinned compute to device 0 regardless of the selected device. Correlate the host-accessible pool to the same id so pinned decode inputs live on the device that runs the model. Single-GPU resolves to id 0, unchanged.
The umbrella EP routes backends by model architecture but OGA never sent it, so every model routed as non-LLM. Emit config.model.type as the ep.amdgpuexecutionprovider.model_arch provider option alongside the existing static-padding hints.
Trailing-comment spacing and argument-continuation alignment flagged by the lint-cpp CI check (clang-format 20.1.0). Formatting only, no behavior change.
Emit the ep.directml.enable_host_accessible provider option so the DirectML backend uses host-accessible decode inputs. Sits alongside the existing static-padding and model_arch config entries.
…ft#2165) Applies the reviewer's suggested wording: the exclusion is about WebGPU, so naming the backends that do zero-init adds nothing.
…(review microsoft#2165) The host-accessible interface was selected in a separate if that reassigned p_device_inputs_ after the chain had already set it. It is now an else if in that chain, so the field is assigned once. The interface getter already returns null when no host-accessible allocator exists, so the extra guard on it went away with the restructure.
…rosoft#2165) EnsureDeviceOrtInit carried the AMD-specific device-id lookup and the host-accessible allocator acquisition inline, including the provider name as a literal. Both now sit behind DeviceInterface virtuals with no-op defaults, following the existing ShapeInitSessionProviderOptions pattern, so the shared path no longer branches on the device: GetDeviceId - id the allocators bind to, 0 unless the device resolves one from EP metadata InitDeviceAllocators - lets a device set up any additional allocators it offers once the device allocator exists The AMDGPU implementations move to src/amdgpu/interface.cpp, which also drops the SetAMDGPUDeviceId free function that existed only to reach back into the interface from the shared path.
…ft#2165) Comments in model.cpp, model.h, smartptrs.h, onnxruntime_api.h and generators.cpp named AMD backends while describing device-agnostic code. The behaviour they document is not vendor-specific, so the names are dropped and the wording shortened.
The reviewer asked for this to be an overload of Create rather than a separately named CreateV2. It turned out to have no callers left: the allocator setup now reads the memory-info the EP advertises instead of building one, so the wrapper is dropped rather than renamed.
…icrosoft#2165) The name sits next to the interface that uses it, so the comment restates what the code already shows.
1e76280 to
05297cc
Compare
A plugin EP is only discoverable once its library is registered on the OrtEnv, and the C model_benchmark has no option to pass a path. Resolve it the way the RyzenAI interface does and register it, skipping entirely when the EP is already registered or the library is not found. Also accept AMDGPU in the benchmark's execution-provider list.
The AMDGPU umbrella integration for OGA lives in upstream microsoft/onnxruntime-genai#2165, which is cut against a base newer than the 0.14 line, so it cannot be applied on the current v0.14.0 pin. Move both workflows to the v0.15.0 tag and patch it with that PR. The PR is a draft that gets rebased in place, so its number no longer identifies its content: a new step resolves each patch PR's head SHA and folds it into the OGA cache key, replacing the hand-bumped mm<N> token that had to be remembered whenever a patch PR moved. 0.15 enables 1DS telemetry by default, which pulls cpp_client_telemetry (plus curl and mbedtls on Linux) into a cold build and would ship Microsoft telemetry inside AMD artifacts, so OGA now builds with --no_telemetry. 0.15's model_benchmark prints "Peak working set size: <n> bytes (<x> GB)" where 0.14 printed "Peak working set size (bytes): <n>", so the OGA benchmark summary and the two perf-report tools would report "-" for peak memory. Match the current format, which is what the perf_test parser in the same workflow already expects. Verified locally on gfx1151 against ORT 1.27: v0.15.0 + pull/2165.patch applies with git am, builds (Ninja, Release), and runs Llama-3.1-8B AWQ int4 through the AMD GPU umbrella at 287 tok/s prefill and 30.0 tok/s decode.
|
Approving this PR. We should update the main README to show that AMD GPU is supported instead of on the roadmap. Looks like there is still a merge conflict to resolve as well. |
|
|
||
| DeviceInterface* p_device_{}; // The device we're running on (matches device_type_) used for things that work the same on all devices | ||
| DeviceInterface* p_device_inputs_{}; // For some model inputs, the device might be the CPU device (all but KV cache currently for WebGPU and DML) | ||
| DeviceInterface* p_device_logits_{}; // Logits are read back on the CPU every step |
There was a problem hiding this comment.
This comment is misleading. Logits can be on the device for cuda and nvtensorrtrtx eps.
|
|
||
| // Inputs-only interface backed by a host-accessible allocation, so the CPU updates the small | ||
| // decode inputs in place with no per-step roundtrip. Null if the device offers no such allocator. | ||
| DeviceInterface* p_host_accessible_inputs = GetAMDGPUPinnedInputsInterface(); |
There was a problem hiding this comment.
The variable reads p_host_accessible_inputs but queries a EP specialized function GetAMDGPUPinnedInputsInterface. Could we move this behind an abstraction such as p_device->GetHostAccessibleDevice()?
Now, that we have so many device interface variables, we could consider moving all behind p_device->GetInputDevice(), p_device->GetLogitsDevice() so the code can have ep specific overrides instead of us having to add EP specialized functions at the top level.
|
A few thoughts:
|
This PR adds the AMDGPU execution provider to ONNX Runtime GenAI. The AMDGPU EP picks a concrete backend at runtime and OGA drives it as a single
DeviceType::AMDGPU, so callers never manage the backend directly. It supersedes the earlier MIGraphX-only framing (#2093 and the previous version of this PR).Users select it as
"amdgpu". We also accept"AMDGPUExecutionProvider", the catalog name the AMD Windows ML EP MSIX advertises, so harnesses that match against WinML-discovered names work without special-casing. Both normalize to"AMDGPU".What's in the three commits
1. AMDGPU EP with GPU-resident KV cache. Adds the EP wiring (
DeviceType::AMDGPU, the dispatch-table entry, provider-name normalization, and thesrc/amdgpu/sources) plus a GPU-residentDeviceInterfacethat keeps the KV cache on the device, removing the per-token CPU/GPU roundtrip. Cross-device copies go through backend-agnostic opaqueDeviceBufferhandles. It also emits the static-shape prefill config (ep.migraphx.static_pad_*andhip_graph_enable) so the backend compiles the prefill once and reuses a captured graph.2. Host-accessible decode inputs. The small decode inputs (input_ids, position_ids, attention_mask) are allocated from a host-accessible allocator that the CPU can write and the GPU can read, so the CPU updates them in place with no per-step copy. The KV cache and scoring stay on the default device interface. If no host-accessible allocator is available, it falls back to the default input path.
3. Logits on CPU. Logits are GPU-written and CPU-read, the opposite of the decode inputs. On AMDGPU those inputs live in a host-accessible heap that is not coherent for CPU reads, so reading logits from it would return stale data. This commit routes logits to the CPU interface through a new
p_logits_member, leaving only the decode inputs pinned.Configuration
or
"AMDGPUExecutionProvider"also works (the catalog form).Design notes
Graph capture is always on for AMDGPU, so OGA sizes the attention mask and KV cache itself. The EP is told to pad the prefill token axis through the static-pad session-config entries, while OGA keeps ownership of mask and KV sizing. The host-accessible and logits routing are both backend-aware and fall back safely to the default path when the allocator isn't there.
Known limitations
Beam search isn't supported (it needs
past_present_share_buffer=true, which needsnum_beams=1). The host-accessible allocator currently assumesdevice_id=0, so multi-GPU is a follow-up.